/* Nate Deisinger (ndeisinger@wisc.edu)
 *
 * This file is part of soot-csi, a project to apply lightweight
 * instrumentation to Java programs to facilitate debugging through
 * the Soot analysis framework.
 *
 * Copyright (c) 2014, Nate Deisinger, University of Wisconsin-Madison
 *   Madison, Wisconsin 53715
 *   All Rights Reserved.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc.,
 *   51 Franklin Street, Fifth Floor,
 *   Boston, MA  02110-1301
 *   USA
 * or retrieve version 2.1 at their website:
 *   http://www.gnu.org/licenses/lgpl-2.1.html
 */
package uwmadison.soot.analysis;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;

import soot.tagkit.SourceLnPosTag;
import soot.tagkit.Tag;
import soot.toolkits.graph.ExceptionalBlockGraph;
import soot.toolkits.graph.ExceptionalUnitGraph;
import soot.toolkits.graph.pdg.EnhancedBlockGraph;
//import soot.toolkits.graph.ExpandedBlockGraph;
import soot.util.*;
import soot.jimple.*;
import soot.*;
import uwmadison.soot.SootCSI;

/**
 * A Soot analysis for adding call coverage instrumentation
 * to a program and outputting required metadata. 
 * 
 * @author Nate Deisinger
 */
public class CallAnalysis {

    private static SootClass callClass, throwableClass;
    protected static SootMethod blankConstructor, dummyMethod,
    nameConstructor, freeMethod, makeTop, getRef, dumpAllTraces;
    protected static SootField localCallArr, ccInUse;
    
    public static HashMap<String, Integer> methodCounts;
    private boolean doDebug = false;

    /**
     * Performs initialization work to link Soot objects to parts of code.
     */
    static
    {
        Scene.v().loadClassAndSupport("uwmadison.soot.analysis.CallTrace");
        callClass = Scene.v().getSootClass("uwmadison.soot.analysis.CallTrace");
        Scene.v().loadClassAndSupport("java.lang.Throwable");
        throwableClass = Scene.v().getSootClass("java.lang.Throwable");
        freeMethod = callClass.getMethod("void freeStack()");
        getRef = callClass.getMethodByName("getRef");
        dumpAllTraces = callClass.getMethodByName("dumpAllTraces");
        
        localCallArr = callClass.getFieldByName("cc_arr");
        ccInUse = callClass.getFieldByName("inUse");
        
        methodCounts = new HashMap<String, Integer>();
    }

    /**
     * A private class to represent a call site as a combination of 
     * line number, Soot Unit, and method signature.
     */
    private class MethodEntry
    {
        private Unit unit;
        private int line;
        private String name;

        public MethodEntry(int line, String name, Unit unit)
        {
            this.line = line;
            this.name = name;
            this.unit = unit;
        }

        public int getLine()
        {
            return line;
        }

        public String getName()
        {
            return name;
        }

        public Unit getUnit()
        {
            return unit;
        }
    }

    /**
     * Prints our call coverage metadata to a given PrintStream.
     */
    private void printCalls(PrintWriter out, String methodSig, List<MethodEntry> methods)
    {
        out.println("#" + methodSig);
        for (int i = 0; i < methods.size(); i++)
        {
            MethodEntry me = methods.get(i);
            out.println(i + "|" + me.getLine() + "|" + me.getName());
        }
    }
    
    /**
     * Checks for traps that catch the first unit and adjusts them
     * to begin after our initialization instrumentation.
     * 
     * Also inserts assignments to get our PathTrace object after
     * the caught-exception identity of a handler.
     */
    private void fixTraps(Body methodBody, Unit safeStart)
    {
        for (Trap t : methodBody.getTraps())
        {
            Unit start = t.getBeginUnit();
            if (start == methodBody.getUnits().getFirst())
            {
                //Move to start after init
                t.setBeginUnit(safeStart);
            }
            Unit handler = t.getHandlerUnit();
            
            //Insert after exception identity
            StaticFieldRef useRef = Jimple.v().newStaticFieldRef(ccInUse.makeRef());
            AssignStmt tempAssign = Jimple.v().newAssignStmt(ccUse, useRef);
            methodBody.getUnits().insertAfter(tempAssign, handler);
        }
    }

    Local myCallArr, myGlobalCallArr, tempBool, ccUse; //Various instrumentation locals we'll add
    
    /**
     * Instruments a function to perform call coverage.
     */
    public CallAnalysis(ExceptionalBlockGraph exceptionalBlockGraph) {
        Body methodBody = exceptionalBlockGraph.getBody();
        ArrayList<MethodEntry> methods = new ArrayList<MethodEntry>();
        SootMethod myMethod = methodBody.getMethod();
        if (myMethod.getDeclaringClass().getName().contains("uwmadison")) return;

        //Walk the method body and find all invoke statements
        //Insert them into a list; length of list = length of array
        //Insert appropriate array + initialization
        //Also insert exception catch + utility class

        //Set up our CallTrace object
        Chain<Local> locals = methodBody.getLocals();
        String ctName = myMethod.getSignature();
        Type ctType = callClass.getType();
        Local myCT = Jimple.v().newLocal(ctName, ctType);

        myCallArr = Jimple.v().newLocal("__cc_arr", ArrayType.v(BooleanType.v(), 1));
        myGlobalCallArr = Jimple.v().newLocal("__cc_global_arr", ArrayType.v(BooleanType.v(), 1));
        tempBool = Jimple.v().newLocal("__cc_temp_bool", BooleanType.v());
        ccUse = Jimple.v().newLocal("__cc_use_bool", BooleanType.v());

        //Add it to the method's locals
        locals.add(myCT);
        locals.add(myCallArr);
        locals.add(myGlobalCallArr);
        locals.add(tempBool);
        locals.add(ccUse);

        int totalCalls = 0;
        Chain<Unit> units = methodBody.getUnits();

        //Find invoke statements
        for (Unit u : units)
        {
            if (u instanceof InvokeStmt)
            {
                InvokeStmt myCall = (InvokeStmt) u;
                if (myCall.containsInvokeExpr())
                {
                    //Get metadata (line number, name) of call
                    if (doDebug) G.v().out.println(myCall);
                    SootMethod calledMethod = myCall.getInvokeExpr().getMethod();
                    SourceLnPosTag lineInfoTag = null;
                    Iterator<Tag> startTagIterator = u.getTags().iterator();
                    while (startTagIterator.hasNext())
                    {
                        Tag tempTag = startTagIterator.next();
                        if (tempTag instanceof SourceLnPosTag)
                        {
                            lineInfoTag = (SourceLnPosTag) tempTag;
                            break;
                        }
                    }
                    int lineNumber;
                    if (lineInfoTag == null) 
                    {
                        lineNumber = -1;
                    }
                    else
                    {
                        lineNumber = lineInfoTag.startLn();
                    }

                    MethodEntry entry = new MethodEntry(lineNumber, calledMethod.getName(), u);
                    methods.add(entry);

                    //Increment total # of calls
                    totalCalls++;
                }
            }
        }

        //Create local/global bools
        Local[] cc_bools = new Local[totalCalls];
        SootField[] global_bools = CallUtility.makeCallFields(myMethod.getSignature(), totalCalls);
        for (int i = 0; i < totalCalls; i++)
        {
            cc_bools[i] = Jimple.v().newLocal("cc_bool_" + i, BooleanType.v());
            methodBody.getLocals().add(cc_bools[i]);
        }

        //Insert instrumentation
        for (int i = 0; i < totalCalls; i++)
        {
            AssignStmt valAssign = Jimple.v().newAssignStmt(cc_bools[i], IntConstant.v(1));
            AssignStmt globalAssign = Jimple.v().newAssignStmt(
                    Jimple.v().newStaticFieldRef(global_bools[i].makeRef()), IntConstant.v(1));
            units.insertBefore(valAssign, methods.get(i).getUnit());
            units.insertBefore(globalAssign, methods.get(i).getUnit());
        }

        //Initialize our local arrays to null
        Unit insertPoint = units.getFirst();
        if (methodBody.getMethod().getParameterCount() > 0)
        {
            //Check if static (eg. no @this-assignment)
            int offset = 0;
            if (methodBody.getMethod().isStatic()) offset = 1;
            //Walk past parameters
            for (int i = offset; i < methodBody.getMethod().getParameterCount(); i++)
            {
                insertPoint = (Stmt) units.getSuccOf(insertPoint);
            }
        }
        else
        {
            insertPoint = units.getFirst();
        }
        
        //For determining trap range
        Unit startTrap = null;
        
        //Initialize our values
        for (int i = 0; i < methods.size(); i++)
        {
            AssignStmt valAssign = Jimple.v().newAssignStmt(cc_bools[i], IntConstant.v(0));
            if (startTrap == null) startTrap = valAssign;
            units.insertAfter(valAssign, insertPoint);
        }
        AssignStmt nullArr = Jimple.v().newAssignStmt(myCallArr, NullConstant.v());
        if (startTrap == null) startTrap = nullArr;
        AssignStmt nullCatch = Jimple.v().newAssignStmt(myCT, NullConstant.v());
        AssignStmt nullGlobalArr = Jimple.v().newAssignStmt(myGlobalCallArr, NullConstant.v());
        AssignStmt tempInit = Jimple.v().newAssignStmt(tempBool, IntConstant.v(0));
        AssignStmt ccUseInit = Jimple.v().newAssignStmt(ccUse, IntConstant.v(0));
        
        //Add instrumentation to method
        units.insertAfter(nullArr, insertPoint);
        units.insertAfter(nullCatch, insertPoint);
        units.insertAfter(nullGlobalArr, insertPoint);
        units.insertAfter(tempInit, insertPoint);
        units.insertAfter(ccUseInit, insertPoint);
        
        //Only trap units after our initialization code
        startTrap = units.getSuccOf(startTrap);

        //Call to factory method
        StaticInvokeExpr constructExpr = Jimple.v().newStaticInvokeExpr(
                getRef.makeRef(), StringConstant.v(myMethod.getSignature()), IntConstant.v(totalCalls));

        //Call to re-init method
        InvokeExpr closeExpr = Jimple.v().newStaticInvokeExpr(freeMethod.makeRef());

        //Finally, insert calls to CallTrace freeStack() before returns
        Iterator stmtIt = units.snapshotIterator();
        while (stmtIt.hasNext())
        {
            Unit u = (Unit) stmtIt.next();
            if ((u instanceof ReturnStmt) || (u instanceof RetStmt)  || (u instanceof ReturnVoidStmt))
            {
                //Quickly check if we have anything in stack that needs freeing
                IfStmt useCheck = Jimple.v().newIfStmt(Jimple.v().newEqExpr(ccUse,IntConstant.v(0)), u);
                units.insertBefore(useCheck, u);
                
                //Insert close op just before
                InvokeStmt closeInvoke = Jimple.v().newInvokeStmt(closeExpr);
                units.insertBefore(closeInvoke, u);
                
                //Ensure proper jump target
                useCheck.setTarget(u);
            }
        }
        
        //Double-check traps don't cover our instrumentation code and instrument catch blocks
        //to check if we're in use
        fixTraps(methodBody, startTrap);

        Unit endStmt = units.getLast();
                
        //Instructions to handle uncaught exception case
        Chain<Unit> instrChain = new HashChain<Unit>();
        
        //Catch uncaught exceptions and store result
        Local catchRefLocal = soot.jimple.Jimple.v().newLocal("e", soot.RefType.v("java.lang.Throwable"));
        methodBody.getLocals().add(catchRefLocal);
        CaughtExceptionRef caughtRef = soot.jimple.Jimple.v().newCaughtExceptionRef();
        Stmt caughtIdentity = soot.jimple.Jimple.v().newIdentityStmt(catchRefLocal, caughtRef);
        instrChain.add(caughtIdentity);

        //Insert call to factory method for CallTrace object
        instrChain.add(Jimple.v().newAssignStmt(myCT, constructExpr));
       
        //Insert assignment to link local/global boolean arrays and ref arrays
        InstanceFieldRef callArrRef = Jimple.v().newInstanceFieldRef(myCT, localCallArr.makeRef());
        AssignStmt linkStmt = Jimple.v().newAssignStmt(myCallArr, callArrRef);
        instrChain.add(linkStmt);
        
        //Insert assignments between local/static bools and the arrays
        for (int i = 0; i < methods.size(); i++)
        {
            //AssignStmt boolAssign = Jimple.v().newAssignStmt(dummyBool, IntConstant.v(1));
            ArrayRef ourArray = Jimple.v().newArrayRef(myCallArr, IntConstant.v(i));
            AssignStmt valAssign = Jimple.v().newAssignStmt(ourArray, cc_bools[i]);
            StaticFieldRef globalRef = Jimple.v().newStaticFieldRef(global_bools[i].makeRef());
            ArrayRef ourGlobalArray = Jimple.v().newArrayRef(myGlobalCallArr, IntConstant.v(i));
            AssignStmt tempAssign = Jimple.v().newAssignStmt(tempBool, globalRef);
            AssignStmt globalValAssign = Jimple.v().newAssignStmt(ourGlobalArray, tempBool);
            instrChain.add(valAssign);
            //instrChain.add(tempAssign);
            //instrChain.add(globalValAssign);
        }

        
        //Check if we're the main method
        if (myMethod.isMain())
        {
            //Special case: dump info
                        StaticInvokeExpr dumpExpr = Jimple.v().newStaticInvokeExpr(
                    dumpAllTraces.makeRef());
            Stmt dumpStmt = Jimple.v().newInvokeStmt(dumpExpr);
            instrChain.add(dumpStmt);
        }
        
        //Throw exception
        ThrowStmt throwStmt = Jimple.v().newThrowStmt(catchRefLocal);
        instrChain.add(throwStmt);
        
        //Insert instrumentation        
        units.insertAfter(instrChain, endStmt);
        
        //Move "endStmt" one forward because newTrap is exclusive on the end of the trap
        endStmt = units.getSuccOf(endStmt);
        
        //Add trap to instrumentation
        methodBody.getTraps().add(Jimple.v().newTrap(throwableClass, startTrap, endStmt, caughtIdentity));

        //Print out metadata info
        String fileName = null;
        try
        {
            fileName = SootCSI.outDir + myMethod.getDeclaringClass().getName() + "_cc";
            PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter( 
                    fileName, true)));
            printCalls(out, myMethod.getSignature(), methods);
            out.close();
        }
        catch (IOException e)
        {
            G.v().out.println("Warning! Could not open file " + fileName);
        }

        if(doDebug)
        {
            G.v().out.println("--------TRANSFORMED--------");
            for (Unit u : exceptionalBlockGraph.getBody().getUnits())
            {
                G.v().out.println(u);
            }
            G.v().out.println("----END TRANSFORMED--------");
        }

    }

}
